{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4227aaea",
   "metadata": {},
   "source": [
    "## 导入模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dcdc21c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sentence_transformers import SentenceTransformer\n",
    "import numpy as np\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c9a4e31",
   "metadata": {},
   "source": [
    "## 导入本地模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c31472b",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = SentenceTransformer(\n",
    "    \"./model/all-MiniLM-L6-v2\", device=\"cuda\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b02f0e08",
   "metadata": {},
   "source": [
    "## 数据文件\n",
    "#### 下载数据集\n",
    "\n",
    "~~~\n",
    "wget https://storage.googleapis.com/generall-shared-data/startups_demo.json\n",
    "~~~\n",
    "\n",
    "#### 读取数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53c3fa40",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.read_json(\"./data/startups_demo.json\", lines=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "04ae882e",
   "metadata": {},
   "source": [
    "## 编码"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ffd0bcd6",
   "metadata": {},
   "outputs": [],
   "source": [
    "vectors = model.encode(\n",
    "    [row.alt + \". \" + row.description for row in df.itertuples()],\n",
    "    show_progress_bar=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99fbac00",
   "metadata": {},
   "outputs": [],
   "source": [
    "    # for idx, row in enumerate(df.itertuples()):\n",
    "    #     print(\"原始文本：\", row.alt + \". \" + row.description)\n",
    "    #     print(\"对应向量：\", vectors[idx])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e2708ffc",
   "metadata": {},
   "source": [
    "## 向量维度"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29d1e1d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "vectors.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc49da97",
   "metadata": {},
   "source": [
    "## 保存向量到npy文件"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c851941",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.save(\"./data/startup_vectors.npy\", vectors, allow_pickle=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3fcccb75",
   "metadata": {},
   "source": [
    "## 在docker中运行Qdrant\n",
    "\n",
    "#### 下载Qdrant镜像\n",
    "\n",
    "~~~\n",
    "docker pull qdrant/qdrant\n",
    "~~~\n",
    "\n",
    "#### 在docker中启动Qdrant\n",
    "\n",
    "~~~\n",
    "docker run -p 6333:6333 \\\n",
    "    -v $(pwd)/qdrant_storage:/qdrant/storage \\\n",
    "    qdrant/qdrant\n",
    "~~~"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9bb3d47d",
   "metadata": {},
   "source": [
    "## 将数据上传到Qdrant"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b2176a1",
   "metadata": {},
   "source": [
    "### 为Qdrant创建一个客户端对象"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89f23b49",
   "metadata": {},
   "outputs": [],
   "source": [
    "from qdrant_client import QdrantClient\n",
    "from qdrant_client.models import VectorParams, Distance\n",
    "import json\n",
    "\n",
    "client = QdrantClient(\"http://localhost:6333\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cfbd53fe",
   "metadata": {},
   "source": [
    "### 创建集合"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d12dd57d",
   "metadata": {},
   "outputs": [],
   "source": [
    "if not client.collection_exists(\"startups\"):\n",
    "    client.create_collection(\n",
    "        collection_name=\"startups\",\n",
    "        vectors_config=VectorParams(size=384, distance=Distance.COSINE),\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e947c3ed",
   "metadata": {},
   "source": [
    "#### Qdrant 客户端库定义了一个特殊函数，允许您将数据集加载到服务中。 但是，由于数据可能太多而无法容纳单个计算机内存，因此该函数将数据的迭代器作为输入。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "989c277a",
   "metadata": {},
   "outputs": [],
   "source": [
    "fd = open(\"./data/startups_demo.json\")\n",
    "\n",
    "# payload is now an iterator over startup data\n",
    "payload = map(json.loads, fd)\n",
    "\n",
    "# Load all vectors into memory, numpy array works as iterable for itself.\n",
    "# Other option would be to use Mmap, if you don't want to load all data into RAM\n",
    "vectors = np.load(\"./data/startup_vectors.npy\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "90c82ae8",
   "metadata": {},
   "source": [
    "### 上传数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b25c895b",
   "metadata": {},
   "outputs": [],
   "source": [
    "client.upload_collection(\n",
    "    collection_name=\"startups\",\n",
    "    vectors=vectors,\n",
    "    payload=payload,\n",
    "    ids=None,  # Vector ids will be assigned automatically\n",
    "    batch_size=256,  # How many vectors will be uploaded in a single request?\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f0cafc1",
   "metadata": {},
   "source": [
    "#### 访问http://localhost:6333/dashboard#/welcome，在collection中可以看到上传的startups数据库"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7bfb455",
   "metadata": {},
   "source": [
    "### 查询"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1eb81327",
   "metadata": {},
   "outputs": [],
   "source": [
    "def search_query(query_text, top_k=5):\n",
    "    # query_vector = model.encode(query_text).tolist()\n",
    "\n",
    "    results = client.search(\n",
    "        collection_name=\"startups\",\n",
    "        query_vector=model.encode(query_text).tolist(),  \n",
    "        limit=top_k,\n",
    "        with_payload=True,\n",
    "    )\n",
    "    return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1fa9b22",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 查询问题\n",
    "query = \"cloud storage security platform\"\n",
    "\n",
    "# 调用 search_query 函数\n",
    "results = search_query(query, top_k=5)\n",
    "\n",
    "# 打印结果\n",
    "print(f\"\\nTop 5 results for: '{query}'\\n\")\n",
    "for i, res in enumerate(results, 1):\n",
    "    alt = res.payload.get(\"alt\", \"N/A\")\n",
    "    desc = res.payload.get(\"description\", \"\")\n",
    "    score = res.score\n",
    "    print(f\"{i}. {alt} (score={score:.4f})\\n   {desc}\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54ee3e97",
   "metadata": {},
   "source": [
    "#### 查看权重文件model.safetensors中的参数信息（与正文无关）"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2177595",
   "metadata": {},
   "outputs": [],
   "source": [
    "from safetensors.torch import safe_open\n",
    "\n",
    "file_path = \"./model/all-MiniLM-L6-v2/model.safetensors\"\n",
    "\n",
    "with safe_open(file_path, framework=\"pt\", device=\"cpu\") as f:\n",
    "    for key in f.keys():\n",
    "        tensor = f.get_tensor(key)\n",
    "        print(f\"{key}: shape={tensor.shape}, dtype={tensor.dtype}\")\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
